#include "compi/compi.h"
#include "compi/context.h"
#include "compi/p2p.h"
#include "compi/tag.h"

#include <gtest/gtest.h>

#include <mpi.h>
#include <string>
#include <vector>

using namespace compi;

class TagLeaseTransportTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        int initialized = 0;
        MPI_Initialized(&initialized);
        if (!initialized)
        {
            int provided;
            MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided);
        }
    }

    void TearDown() override { Context::destroy_for_comm(MPI_COMM_WORLD); }
};

TEST_F(TagLeaseTransportTest, SerializedSendLeaseKeepsOwnership)
{
    auto & ctx = Context::for_comm(MPI_COMM_WORLD);
    auto & alloc = ctx.tag_allocator();
    auto ns = alloc.create_namespace("tl_serialized");

    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    if (rank == 0)
    {
        TagLease lease = ns.lease();
        const std::string payload = "lease-payload";

        auto req = isend_serialized(ctx, /*dest=*/1, payload, lease);
        EXPECT_TRUE(req.wait().is_success());

        // 完成后标签仍由租约持有：新请求不应复用同一标签值
        auto next_tag = ns.request();
        EXPECT_NE(next_tag.value, lease.value());
        ns.release(next_tag);

        // 手动归还后应复用原标签（LIFO）
        lease.retire();
        auto reused = ns.request();
        EXPECT_EQ(reused.value, lease.value());
        ns.release(reused);
    }
    else if (rank == 1)
    {
        std::vector<std::byte> recv;
        auto rreq = irecv_serialized(ctx, /*src=*/0, recv, ns);
        ASSERT_TRUE(rreq.wait().is_success());

        std::string got(reinterpret_cast<const char *>(recv.data()), recv.size());
        EXPECT_EQ(got, "lease-payload");
    }
}

// 自定义 main：统一初始化与析构，避免 mpiexec 关于未调用 Finalize 的异常退出
static compi::Environment * g_test_env = nullptr;

int
main(int argc, char ** argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    g_test_env = new compi::Environment();
    const int code = RUN_ALL_TESTS();
    delete g_test_env;
    g_test_env = nullptr;
    return code;
}

TEST_F(TagLeaseTransportTest, FixedSendLeaseKeepsOwnership)
{
    auto & ctx = Context::for_comm(MPI_COMM_WORLD);
    auto & alloc = ctx.tag_allocator();
    auto ns = alloc.create_namespace("tl_fixed");

    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    if (rank == 0)
    {
        TagLease lease = ns.lease();
        int value = 1234;

        auto req = isend_fixed(ctx, /*dest=*/1, &value, MPI_INT, lease);
        EXPECT_TRUE(req.wait().is_success());

        // 完成后标签仍由租约持有：新请求不应复用同一标签值
        auto next_tag = ns.request();
        EXPECT_NE(next_tag.value, lease.value());
        ns.release(next_tag);

        // 手动归还后应复用原标签（LIFO）
        lease.retire();
        auto reused = ns.request();
        EXPECT_EQ(reused.value, lease.value());
        ns.release(reused);
    }
    else if (rank == 1)
    {
        int recv = 0;
        auto rreq = irecv_fixed(ctx, /*src=*/0, &recv, MPI_INT, ns);
        ASSERT_TRUE(rreq.wait().is_success());
        EXPECT_EQ(recv, 1234);
    }
}